 /*
  * Copyright (c) 2018. paascloud.net All Rights Reserved.
  * 项目名称：paascloud快速搭建企业级分布式微服务平台
  * 类名称：TokenJwtEnhancer.java
  * 创建人：刘兆明
  * 联系方式：paascloud.net@gmail.com
  * 开源地址: https://github.com/paascloud
  * 博客地址: http://blog.paascloud.net
  * 项目官网: http://paascloud.net
  */

 package com.weixing.oauth2.common.provider.token;

 import com.fasterxml.jackson.core.JsonProcessingException;
 import com.fasterxml.jackson.databind.ObjectMapper;
 import com.weixing.mall.base.utils.AuthCryptoUtil;
 import com.weixing.oauth2.common.properties.OAuth2Properties;
 import com.weixing.oauth2.common.userdetails.SecurityUser;
 import org.springframework.security.core.Authentication;
 import org.springframework.security.core.userdetails.UserDetails;
 import org.springframework.security.oauth2.common.DefaultOAuth2AccessToken;
 import org.springframework.security.oauth2.common.OAuth2AccessToken;
 import org.springframework.security.oauth2.provider.OAuth2Authentication;
 import org.springframework.security.oauth2.provider.token.TokenEnhancer;

 import java.util.HashMap;
 import java.util.Map;


 /**
  * The class Token jwt enhancer.
  *
  * @author paascloud.net @gmail.com
  */
 public class TokenJwtEnhancer implements TokenEnhancer {

     private ObjectMapper objectWrapper = new ObjectMapper();
     private OAuth2Properties auth2Properties;

     public TokenJwtEnhancer(OAuth2Properties auth2Properties) {
         this.auth2Properties = auth2Properties;
     }

     /**
      * Enhance o auth 2 access token.
      *
      * @param accessToken          the access token
      * @param oAuth2Authentication the o auth 2 authentication
      *
      * @return the o auth 2 access token
      */
     @Override
     public OAuth2AccessToken enhance(OAuth2AccessToken accessToken, OAuth2Authentication oAuth2Authentication) {
         Map<String, Object> info = new HashMap<>(8);
         info.put("timestamp", System.currentTimeMillis());
         Authentication authentication = oAuth2Authentication.getUserAuthentication();
         if (authentication != null) {
             Object principal = authentication.getPrincipal();
             if (principal instanceof UserDetails) {
                 SecurityUser authUser = (SecurityUser) principal;
                 String userRole = authUser.getUserRole();
                 info.put("ur", userRole);
                 Map<String, Object> extData = authUser.getExtData();
                 if(null != extData){
                     for (String key : extData.keySet()) {
                         info.put(key, extData.get(key));
                     }
                 }

                 Object sensitiveData = authUser.getSensitiveData();
                 String senstitiveJsonStr = "";
                 try {
                     senstitiveJsonStr = objectWrapper.writeValueAsString(sensitiveData);
                 } catch (JsonProcessingException e) {
                     e.printStackTrace();
                 }
                 info.put("uk", AuthCryptoUtil.encrypt(senstitiveJsonStr));
             }
         }

         ((DefaultOAuth2AccessToken) accessToken).setAdditionalInformation(info);

         return accessToken;
     }
 }
